//
// Created by wdq on 2022/9/11.
//
#include "iostream"
#include <cstring>
#include <map>

struct HuffmanNode {
    char data;
    float weight;
    char *code;
    int left, right;
    int parent;

    HuffmanNode() {
        left = right = parent = 0;
        code = nullptr;

    }

    ~HuffmanNode() {
        delete code;
    }

    void output() const {
        printf("%.f\t%d\t%d\t%d\t", weight, left, right, parent);
        if (code != nullptr) {
            printf("%c\t%s", data, code);
        }
        printf("\n");
    }

    bool isLeaf() const {
        //千万不要 left == right == 0，这样写因为left会和right先比较结果就为1，1==0为假
        return left == 0 && right == 0;
    }

};

struct HuffmanTree {
    int nums;
    HuffmanNode *nodes;
    bool *isChoose;

    HuffmanTree(int nums) {
        this->nums = nums;
        nodes = new HuffmanNode[2 * nums];
        isChoose = new bool[2 * nums];
        memset(isChoose, 0, 2 * nums);

    }

    ~HuffmanTree() {
        delete[]nodes;
        delete[]isChoose;
    }

    void createTree(char *str, const int weight[]) {
        memset(isChoose,0,2*nums);
        int length = int(strlen(str));

        for (int i = 1; i <= length; i++) {
            nodes[i] = HuffmanNode();
            nodes[i].data = str[i - 1];
            nodes[i].weight = float(weight[i - 1]);
        }
        for (int i = length + 1; i < 2 * nums; i++) {
            int idx1, idx2;
            findTwoMin(i, idx1, idx2);
            nodes[i] = HuffmanNode();
            nodes[i].weight = nodes[idx1].weight + nodes[idx2].weight;
            nodes[i].left = idx1;
            nodes[i].right = idx2;
            nodes[idx1].parent = i;
            nodes[idx2].parent = i;
        }
        char code[10] = {};
        fillCode(2 * nums - 1, code, 0);

    }

    void findTwoMin(int i, int &idx1, int &idx2) {
        idx1 = findMin(i);
        idx2 = findMin(i);

    }

    int findMin(int i) {
        int length = 0;
        auto *temp = new HuffmanNode[2 * nums];
        int *index = new int[2 * nums];

        for (int j = 1; j < i; j++) {
            if (!isChoose[j]) {
                temp[length] = nodes[j];
                index[length] = j;
                length++;
            }
        }
        int result = 0;
        for (int j = 1; j < length; j++) {
            if (temp[result].weight > temp[j].weight) {
                result = j;
            }
        }
        isChoose[index[result]] = true;
        return index[result];

    }

    void fillCode(int i, char *code, int len) {
        HuffmanNode &node = nodes[i];
        if (node.left != 0) {
            code[len] = '0';
            fillCode(node.left, code, len + 1);
        }
        if (node.right != 0) {
            code[len] = '1';
            fillCode(node.right, code, len + 1);
        }
        if (node.isLeaf()) {
            if (node.code == nullptr) {
                node.code = new char[10];
            }
            strcpy(node.code, code);
        }

    }

    void output() const {
        printf("nums\tweight\tleft\tright\tparent\tchar\tcode\n");
        for (int i = 1; i < 2 * nums; i++) {
            printf("%d\t", i);
            nodes[i].output();
        }
    }


    void encode(char *src, char *dst) {
        std::map<char, char *> res;
        for (int i = 1; i < 2 * nums; i++) {
            if (nodes[i].isLeaf()) {
                res.insert(std::pair<char, char *>(nodes[i].data, nodes[i].code));
            }
        }
        int length = 0;
        for (int i = 0; i < strlen(src); i++) {
            auto code = res.at(src[i]);
            strcpy(dst + length, code);
            length += int(strlen(code));

        }
    }

    void decode(char *src, char *dst) {
        int length = 0;
        dfs(src, dst, nodes[2 * nums - 1], length, 0);


    }

    void dfs(char *src, char *dst, const HuffmanNode &node, int &length, int srcLength) {
        if (node.isLeaf()) {
            dst[length++] = node.data;
            dfs(src, dst, nodes[2 * nums - 1], length, srcLength);
        } else if (int(strlen(src) <= srcLength)) {
            return;
        } else if (src[srcLength] == '0') {
            dfs(src, dst, nodes[node.left], length, srcLength + 1);
        } else if (src[srcLength] == '1') {
            dfs(src, dst, nodes[node.right], length, srcLength + 1);
        }


    }
};

int main() {
    HuffmanTree tree(4);
    char *str = "abcd";
    int weight[] = {7, 5, 2, 4};
    tree.createTree(str, weight);
    tree.output();

    char *p = new char[100];
    memset(p, 0, 100);
    tree.encode("abb", p);
    printf("abb encode result: %s\n", p);

    char *q = new char[100];
    memset(q, 0, 100);
    tree.decode("1110111", q);
    printf("1110111 decode result: %s\n", q);
    return 0;
}